# -*- coding: utf-8 -*-
"""
@Time ： 2025/2/20 11:12
@Auth ： 九问
@File ：crypto_tool.py
@IDE ：PyCharm
@Email : 2750826557@qq.com
"""
import binascii
import base64
from Crypto.Cipher import DES, AES  # 需使用pycryptodome库
from Crypto.Util.Padding import pad, unpad
from Crypto.Hash import SHA256


class MyCrypto:
    def __init__(self, key, iv, mode='DES'):
        """
        初始化加密类
        :param key: 加密密钥（字节类型）
        :param iv: 初始化向量（字节类型）
        :param mode: 加密模式，'DES' 或 'AES'（默认值为 'DES'）
        """
        self.mode = mode
        self.key = self._process_key(key)
        self.iv = self._process_iv(iv)

    def _process_key(self, key):
        """
        处理密钥，确保其长度符合加密算法的要求
        :param key: 原始密钥（字节类型）
        :return: 处理后的密钥（字节类型）
        """
        if self.mode == 'DES':
            # DES 密钥长度必须为 8 字节
            return self._adjust_length(key, 8)
        elif self.mode == 'AES':
            # AES 密钥长度必须为 16、24 或 32 字节
            # 使用 SHA-256 哈希函数处理密钥，确保其长度为 32 字节
            hasher = SHA256.new(key)
            return hasher.digest()
        else:
            raise ValueError("Unsupported encryption mode")

    def _process_iv(self, iv):
        """
        处理初始向量，确保其长度符合加密算法的要求
        :param iv: 原始初始向量（字节类型）
        :return: 处理后的初始向量（字节类型）
        """
        if self.mode == 'DES':
            # DES IV 长度必须为 8 字节
            return self._adjust_length(iv, 8)
        elif self.mode == 'AES':
            # AES IV 长度必须为 16 字节
            return self._adjust_length(iv, 16)
        else:
            raise ValueError("Unsupported encryption mode")

    def _adjust_length(self, data, required_length):
        """
        调整数据的长度，通过截断或填充的方式
        :param data: 原始数据（字节类型）
        :param required_length: 需要的长度
        :return: 调整后的数据（字节类型）
        """
        if len(data) > required_length:
            return data[:required_length]
        elif len(data) < required_length:
            return data + b'\x00' * (required_length - len(data))
        else:
            return data

    def encrypt(self, data):
        """
        加密数据
        :param data: 要加密的数据（字符串类型）
        :return: 加密后的数据（十六进制字符串或 Base64 编码字符串）
        """
        data = data.encode()
        if self.mode == 'DES':
            cipher = DES.new(self.key, DES.MODE_CFB, self.iv, segment_size=64)  # 设置 CFB 模式的块大小为 8 字节
            encrypted_data = cipher.encrypt(data)
            return binascii.b2a_hex(encrypted_data).decode("utf-8")
        elif self.mode == 'AES':
            cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
            encrypted_data = cipher.encrypt(pad(data, AES.block_size))
            return base64.b64encode(encrypted_data).decode("utf-8")
        else:
            raise ValueError("Unsupported encryption mode")

    def decrypt(self, data):
        """
        解密数据
        :param data: 要解密的数据（十六进制字符串或 Base64 编码字符串）
        :return: 解密后的数据（字符串）
        """
        if self.mode == 'DES':
            data = binascii.a2b_hex(data.encode(encoding='utf-8'))
            cipher = DES.new(self.key, DES.MODE_CFB, self.iv, segment_size=64)  # 设置 CFB 模式的块大小为 8 字节
            decrypted_data = cipher.decrypt(data)
            return decrypted_data.decode()
        elif self.mode == 'AES':
            data = base64.b64decode(data)
            cipher = AES.new(self.key, AES.MODE_CBC, self.iv)
            decrypted_data = unpad(cipher.decrypt(data), AES.block_size)
            return decrypted_data.decode()
        else:
            raise ValueError("Unsupported decryption mode")


myDESCipher = MyCrypto(b'jiuwen', b'12345678', mode='DES')
myAESCipher = MyCrypto(b'jiuwen', b'1234567890123456', mode='AES')
# 示例用法
if __name__ == '__main__':
    # DES 示例
    myDESCipher = MyCrypto(b'jiuwen', b'12345678', mode='DES')
    encrypted = myDESCipher.encrypt('qwe123')
    print("DES 加密后的数据:", encrypted)
    decrypted = myDESCipher.decrypt(encrypted)
    print("DES 解密后的数据:", decrypted)

    # AES 示例
    myAESCipher = MyCrypto(b'jiuwen', b'1234567890123456', mode='AES')
    encrypted = myAESCipher.encrypt('qwe123')
    print("AES 加密后的数据:", encrypted)
    decrypted = myAESCipher.decrypt(encrypted)
    print("AES 解密后的数据:", decrypted)
